--- A 'Map' based on a balanced binary tree 
package frege.data.Map inline (!!) where

import frege.Prelude hiding (toList, join, filter, map, null, empty, !!)

type Size = Int
data Map k a = Tip | Bin Size k a (Map k a) (Map k a)

fromList :: (Ord k) => [(k, v)] -> Map k v
fromList xs = fold ins empty xs where
    ins t (k,x) = insert k x t

empty :: Map k a
empty = Tip

insert :: Ord k => k -> a -> Map k a -> Map k a
insert = go
  where
    go kx x Tip = singleton kx x
    go kx x (Bin sz ky y l r) =
        case compare kx ky of
            LT -> balanceL ky y (go kx x l) r
            GT -> balanceR ky y l (go kx x r)
            EQ -> Bin sz kx x l r

null :: Map k a -> Bool
null Tip      = true
null (Bin {}) = false

size :: Map k a -> Int
size Tip              = 0
size (Bin sz _ _ _ _) = sz

singleton :: k -> a -> Map k a
singleton k x = Bin 1 k x Tip Tip

lookup :: Ord k => k -> Map k a -> Maybe a
lookup _ Tip = Nothing
lookup k (Bin _ kx x l r) =
        case compare k kx of
            LT -> lookup k l
            GT -> lookup k r
            EQ -> Just x

member :: Ord k => k -> Map k a -> Bool
member k m = case lookup k m of
    Nothing -> false
    Just _  -> true

map :: (a -> b) -> Map k a -> Map k b
map f = mapWithKey (\_ -> (\x -> f x))

mapWithKey :: (k -> a -> b) -> Map k a -> Map k b
mapWithKey _ Tip = Tip
mapWithKey f (Bin sx kx x l r) = Bin sx kx (f kx x) (mapWithKey f l) (mapWithKey f r)

filter :: Ord k => (a -> Bool) -> Map k a -> Map k a
filter p m
  = filterWithKey (\_ -> (\x -> p x)) m

filterWithKey :: Ord k => (k -> a -> Bool) -> Map k a -> Map k a
filterWithKey _ Tip = Tip
filterWithKey p (Bin _ kx x l r)
  | p kx x    = join kx x (filterWithKey p l) (filterWithKey p r)
  | otherwise = merge (filterWithKey p l) (filterWithKey p r)

toList :: Map k a -> [(k,a)]
toList t = toAscList t

toAscList :: Map k a -> [(k,a)]
toAscList t = foldrWithKey (\k -> (\x -> (\xs -> (k,x):xs))) [] t

toDescList :: Map k a -> [(k,a)]
toDescList t = foldlWithKey (\xs -> (\k -> (\x -> (k,x):xs))) [] t

keys  :: Map k a -> [k]
keys m
  = [k | (k,_) <- assocs m]

assocs :: Map k a -> [(k,a)]
assocs m
  = toList m

elems :: Map k a -> [a]
elems m
  = [x | (_,x) <- assocs m]

fromListWith :: Ord k => (a -> a -> a) -> [(k,a)] -> Map k a 
fromListWith f xs
  = fromListWithKey (\_ -> (\x -> (\y -> f x y))) xs

fromListWithKey :: Ord k => (k -> a -> a -> a) -> [(k,a)] -> Map k a 
fromListWithKey f xs 
  = foldlStrict ins empty xs
  where
    ins t (k,x) = insertWithKey f k x t

insertWithKey :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a
insertWithKey = go
  where
    go _ kx x Tip = singleton kx x
    go f kx x (Bin sy ky y l r) =
        case compare kx ky of
            LT -> balanceL ky y (go f kx x l) r
            GT -> balanceR ky y l (go f kx x r)
            EQ -> Bin sy kx (f kx x y) l r

foldlStrict :: (a -> b -> a) -> a -> [b] -> a
foldlStrict f = go
  where
    go z []     = z
    go z (x:xs) = let z' = f z x in z' `seq` go z' xs

insertWith :: Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
insertWith f = insertWithKey (\_ -> (\x' -> (\y' -> f x' y')))

(!!) :: Ord k => Map k a -> k -> a
m !! k = find k m

find :: Ord k => k -> Map k a -> a
find k m = case lookup k m of
    Nothing -> error "Map.find: element not in the map"
    Just x  -> x


join :: Ord k => k -> a -> Map k a -> Map k a -> Map k a
join kx x Tip r  = insertMin kx x r
join kx x l Tip  = insertMax kx x l
join kx x (l@(Bin sizeL ky y ly ry)) (r@(Bin sizeR kz z lz rz))
  | delta*sizeL < sizeR  = balanceL kz z (join kx x l lz) rz
  | delta*sizeR < sizeL  = balanceR ky y ly (join kx x ry r)
  | otherwise            = bin kx x l r

bin :: k -> a -> Map k a -> Map k a -> Map k a
bin k x l r
  = Bin (size l + size r + 1) k x l r

merge :: Map k a -> Map k a -> Map k a
merge Tip r   = r
merge l Tip   = l
merge (l@(Bin sizeL kx x lx rx)) (r@(Bin sizeR ky y ly ry))
  | delta*sizeL < sizeR = balanceL ky y (merge l ly) ry
  | delta*sizeR < sizeL = balanceR kx x lx (merge rx r)
  | otherwise           = glue l r

glue :: Map k a -> Map k a -> Map k a
glue Tip r = r
glue l Tip = l
glue l r   
  | size l > size r = let ((km,m),l') = deleteFindMax l in balanceR km m l' r
  | otherwise       = let ((km,m),r') = deleteFindMin r in balanceL km m l r'

deleteFindMin :: Map k a -> ((k,a),Map k a)
deleteFindMin t 
  = case t of
      Bin _ k x Tip r -> ((k,x),r)
      Bin _ k x l r   -> let (km,l') = deleteFindMin l in (km,balanceR k x l' r)
      Tip             -> (error "Map.deleteFindMin: can not return the minimal element of an empty map", Tip)

deleteFindMax :: Map k a -> ((k,a),Map k a)
deleteFindMax t
  = case t of
      Bin _ k x l Tip -> ((k,x),l)
      Bin _ k x l r   -> let (km,r') = deleteFindMax r in (km,balanceL k x l r')
      Tip             -> (error "Map.deleteFindMax: can not return the maximal element of an empty map", Tip)

insertMax,insertMin :: k -> a -> Map k a -> Map k a 
insertMax kx x t
  = case t of
      Tip -> singleton kx x
      Bin _ ky y l r
          -> balanceR ky y l (insertMax kx x r)

insertMin kx x t
  = case t of
      Tip -> singleton kx x
      Bin _ ky y l r
          -> balanceL ky y (insertMin kx x l) r

delta,ratio :: Int
delta = 3
ratio = 2

foldlWithKey :: (a -> k -> b -> a) -> a -> Map k b -> a
foldlWithKey f = go
  where
    go z Tip              = z
    go z (Bin _ kx x l r) = go (f (go z l) kx x) r

foldrWithKey :: (k -> a -> b -> b) -> b -> Map k a -> b
foldrWithKey f = go
  where
    go z Tip             = z
    go z (Bin _ kx x l r) = go (f kx x (go z r)) l

balanceL :: k -> a -> Map k a -> Map k a -> Map k a
balanceL k x l r = case r of
  Tip -> case l of
           Tip -> Bin 1 k x Tip Tip
           (Bin _ _ _ Tip Tip) -> Bin 2 k x l Tip
           (Bin _ lk lx Tip (Bin _ lrk lrx _ _)) -> Bin 3 lrk lrx (Bin 1 lk lx Tip Tip) (Bin 1 k x Tip Tip)
           (Bin _ lk lx (ll@(Bin _ _ _ _ _)) Tip) -> Bin 3 lk lx ll (Bin 1 k x Tip Tip)
           (Bin ls lk lx (ll@(Bin lls _ _ _ _)) (lr@(Bin lrs lrk lrx lrl lrr)))
             | lrs < ratio*lls -> Bin (1+ls) lk lx ll (Bin (1+lrs) k x lr Tip)
             | otherwise -> Bin (1+ls) lrk lrx (Bin (1+lls+size lrl) lk lx ll lrl) (Bin (1+size lrr) k x lrr Tip)

  (Bin rs _ _ _ _) -> case l of
           Tip -> Bin (1+rs) k x Tip r

           (Bin ls lk lx ll lr)
              | ls > delta*rs  -> case (ll, lr) of
                   (Bin lls _ _ _ _, Bin lrs lrk lrx lrl lrr)
                     | lrs < ratio*lls -> Bin (1+ls+rs) lk lx ll (Bin (1+rs+lrs) k x lr r)
                     | otherwise -> Bin (1+ls+rs) lrk lrx (Bin (1+lls+size lrl) lk lx ll lrl) (Bin (1+rs+size lrr) k x lrr r)
                   (_, _) -> error "Failure in Data.Map.balanceL"
              | otherwise -> Bin (1+ls+rs) k x l r

balanceR :: k -> a -> Map k a -> Map k a -> Map k a
balanceR k x l r = case l of
  Tip -> case r of
           Tip -> Bin 1 k x Tip Tip
           (Bin _ _ _ Tip Tip) -> Bin 2 k x Tip r
           (Bin _ rk rx Tip (rr@(Bin _ _ _ _ _))) -> Bin 3 rk rx (Bin 1 k x Tip Tip) rr
           (Bin _ rk rx (Bin _ rlk rlx _ _) Tip) -> Bin 3 rlk rlx (Bin 1 k x Tip Tip) (Bin 1 rk rx Tip Tip)
           (Bin rs rk rx (rl@(Bin rls rlk rlx rll rlr)) (rr@(Bin rrs _ _ _ _)))
             | rls < ratio*rrs -> Bin (1+rs) rk rx (Bin (1+rls) k x Tip rl) rr
             | otherwise -> Bin (1+rs) rlk rlx (Bin (1+size rll) k x Tip rll) (Bin (1+rrs+size rlr) rk rx rlr rr)

  (Bin ls _ _ _ _) -> case r of
           Tip -> Bin (1+ls) k x l Tip

           (Bin rs rk rx rl rr)
              | rs > delta*ls  -> case (rl, rr) of
                   (Bin rls rlk rlx rll rlr, Bin rrs _ _ _ _)
                     | rls < ratio*rrs -> Bin (1+ls+rs) rk rx (Bin (1+ls+rls) k x l rl) rr
                     | otherwise -> Bin (1+ls+rs) rlk rlx (Bin (1+ls+size rll) k x l rll) (Bin (1+rrs+size rlr) rk rx rlr rr)
                   (_, _) -> error "Failure in Data.Map.balanceR"
              | otherwise -> Bin (1+ls+rs) k x l r

              